// KdTreeVisualizer.cpp : Defines the entry point for the console application.
//
#include "../RayTraceGPU/MeshGeometry.h"
#include "../RayTraceGPU/ObjModel.h"
#include <Windows.h>
#include "Basic/GxLibBasic.h"
#include "Windows/GxWindows.h"
#include "GL/gl.h"
#include "GL/glu.h"
#include "KdDebugTracer.h"
#include <set>
using namespace GxWinGUI;
using namespace GxLibBasic;
using namespace RayTracePro;

const char* ModelFileName = "../Demo/Media/sponza.obj";

class MainForm : public GLForm
{
private:
	RefPtr<Timer> timer;
	ObjModel mdl;
	RayTraceSystem * system;
	RefPtr<MeshGeometry> mesh;
	KdTree<MeshFace> * tree;
	Array<BBox, true> nodeBounds;
	Array<int, true> selectedNodes;
	Array<int, true> tracedNodes;
	Vec3 viewCenter;
	float viewTheta, viewPhi, viewRadius;
	bool onAdjustView;
	int lastX, lastY;
	int selectedFace;
	int highlightStep;
	Ray testRay;
	void InitModel()
	{
		LoadObj(mdl, ModelFileName);
		RecomputeNormals(mdl);
		mesh = new MeshGeometry();
		MeshGeometry & geometry = *mesh;
		geometry.Normals.Add(mdl.Normals);
		geometry.TexCoords.Add(mdl.TexCoords);
		Vec3 verts[3];
		for (int i = 0; i < mdl.Faces.Count(); i++)
		{
			ObjFace & face = mdl.Faces[i];
			verts[0] = mdl.Vertices[face.VertexIds[0]];
			verts[1] = mdl.Vertices[face.VertexIds[1]];
			verts[2] = mdl.Vertices[face.VertexIds[2]];
			geometry.AddFace(verts, face.NormalIds, face.TexCoordIds, -1);
		}
		String kdFileName = String(ModelFileName) + L".kd";
		if (!FileExists(kdFileName))
		{
			geometry.UpdateAccelerateStructure();
			geometry.SaveKdTree(kdFileName);
		}
		else
		{
			geometry.LoadKdTree(kdFileName);
		}
		tree = mesh->GetKdTree();
		nodeBounds.SetSize(tree->Nodes.Count());
		
		EvalNodeBounds(tree->Bounds, 0);
		Ray ray;
			
		ray.Origin = Vec3(-1356.888062f, 642.013306f, -115.602097f);
		ray.Direction = Vec3(0.818294f, 0.080853f, 0.569086f);//(0.820417, -0.029496, 0.571004)
		ray.ReciprocalDirection.x = 1.0f/ray.Direction.x;
		ray.ReciprocalDirection.y = 1.0f/ray.Direction.y;
		ray.ReciprocalDirection.z = 1.0f/ray.Direction.z;
		ray.tMin = 1.0f;
		ray.tMax = 100000.0f;
		SetTestRay(ray);
	}

	void SetTestRay(const Ray & ray)
	{
		testRay = ray;
		tracedNodes.Clear();

		MeshTracer<false> tracer;
		tracer.mesh = mesh.operator->();
		DifferentialGeometry inter;
	
		RayDifferential rayDiff;
		memset(&rayDiff, 0, sizeof(RayDifferential));
		float t;
		bool intersect = TraverseKdTreeStacklessDbg<MeshFace, MeshTracer<false>, false, false>(system, tracer, inter, *tree, testRay, rayDiff, t, tracedNodes);
		if (intersect)
			printf("Intersection (%f, %f, %f) normal (%f, %f, %f)\n", inter.Position.x, inter.Position.y, inter.Position.z, inter.Normal.x, inter.Normal.y, inter.Normal.z);
		else
		{
			printf("Ray doesnot intersect\n");
		}
	}

	void EvalNodeBounds(BBox & bounds, int id)
	{
		nodeBounds[id] = bounds;
		if (!tree->Nodes[id].GetIsLeaf())
		{
			int dim = tree->Nodes[id].GetSplitDimension();
			BBox bounds1, bounds2;
			bounds1 = bounds;
			bounds2 = bounds;
			if (dim == 0)
			{
				bounds1.xMax = bounds2.xMin = tree->Nodes[id].SplitPlane;
			}
			else if (dim == 1)
			{
				bounds1.yMax = bounds2.yMin = tree->Nodes[id].SplitPlane;
			}
			else
			{
				bounds1.zMax = bounds2.zMin = tree->Nodes[id].SplitPlane;
			}
			EvalNodeBounds(bounds1, id + 1);
			EvalNodeBounds(bounds2, (tree->Nodes[id].GetOffset()>>3));
		}
	}

	void SetProjection()
	{
		glMatrixMode(GL_PROJECTION);
		glLoadIdentity();
		gluPerspective(60.0f, GetClientWidth()/(float)GetClientHeight(), 1.0f, 10000.0f);
		glMatrixMode(GL_MODELVIEW);
	}

	void InitRendering()
	{
		viewTheta = viewPhi = 0.0f;
		viewCenter.SetZero();
		viewCenter.y = 8.0f;
		viewRadius = 30.0f;
		glEnable(GL_DEPTH_TEST);
		glEnable(GL_LIGHTING);
		glEnable(GL_LIGHT0);
		glEnable(GL_COLOR_MATERIAL);
		glColorMaterial(GL_FRONT_AND_BACK, GL_AMBIENT_AND_DIFFUSE);
		float position[4] = {1.0f, 5.0f, 2.0f, 0.0f};
		float lightDiffuse[4] = {0.6f, 0.6f, 0.6f, 1.0f};
		float lightAmbient[4] = {0.3f, 0.3f, 0.3f, 1.0f};
		float lightSpecular[4] = {0.2f, 0.2f, 0.2f, 1.0f};
		glLightfv(GL_LIGHT0, GL_POSITION, position);
		glLightfv(GL_LIGHT0, GL_AMBIENT, lightAmbient);
		glLightfv(GL_LIGHT0, GL_DIFFUSE, lightDiffuse);
		glLightfv(GL_LIGHT0, GL_SPECULAR, lightSpecular);
		SetProjection();
		glEnable(GL_BLEND);
		glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
	}
	std::set<int> highlightFaces;
	void DrawModel()
	{
		glBegin(GL_TRIANGLES);
		for (int i = 0; i<mdl.Faces.Count(); i++)
		{
			if (i == selectedFace)
				glColor3ub(255, 0, 0);
			else if (highlightFaces.find(i) != highlightFaces.end())
				glColor3ub(0,255,0);
			else
				glColor3ub(255, 255, 255);
			glNormal3fv((float*)&mdl.Normals[mdl.Faces[i].NormalIds[0]]);
			glVertex3fv((float*)&mdl.Vertices[mdl.Faces[i].VertexIds[0]]);
			glNormal3fv((float*)&mdl.Normals[mdl.Faces[i].NormalIds[1]]);
			glVertex3fv((float*)&mdl.Vertices[mdl.Faces[i].VertexIds[1]]);
			glNormal3fv((float*)&mdl.Normals[mdl.Faces[i].NormalIds[2]]);
			glVertex3fv((float*)&mdl.Vertices[mdl.Faces[i].VertexIds[2]]);
			
		}
		glEnd();
	}

	union ColorStruct
	{
		struct
		{
			unsigned char R,G,B,A;
		};
		int Value;
	};
	
	int PickFace(int x, int y)
	{
		BeginGLDrawing();
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
		UpdateViewTransform();
		glDisable(GL_LIGHTING);
		glBegin(GL_TRIANGLES);
		for (int i = 0; i<mdl.Faces.Count(); i++)
		{
			ColorStruct c;
			c.Value = i;
			glColor3ub(c.R, c.G, c.B);
			glVertex3fv((float*)&mdl.Vertices[mdl.Faces[i].VertexIds[0]]);
			glVertex3fv((float*)&mdl.Vertices[mdl.Faces[i].VertexIds[1]]);
			glVertex3fv((float*)&mdl.Vertices[mdl.Faces[i].VertexIds[2]]);	
		}
		glEnd();
		
		ColorStruct c;
		glReadPixels(x, GetClientHeight()-y, 1, 1, GL_RGBA, GL_UNSIGNED_BYTE, &c);
		c.A = 0;
		glEnable(GL_LIGHTING);
		EndGLDrawing();
		return c.Value;
	}

	void MainForm_MouseDown(Object * sender, MouseEventArgs args)
	{
		if (args.Button == GxWinGUI::mbRight)
		{
			selectedFace = PickFace(args.X, args.Y);
			printf("PickFace(%d, %d) = %d\n", args.X, args.Y, selectedFace);
			selectedNodes.Clear();
			if (selectedFace > 0)
			{
				viewCenter = mdl.Vertices[mdl.Faces[selectedFace].VertexIds[0]];
			}
			for (int i = 0; i<tree->Nodes.Count(); i++)
			{
				if (tree->Nodes[i].GetIsLeaf())
				{
					int c = tree->Nodes[i].GetElementCount();
					for (int j = 0; j<(int)tree->Nodes[i].GetElementCount(); j++)
					{
						int * elems = (int *)((char*)tree->Leaves.Buffer() + tree->Nodes[i].LeafPtr + sizeof(KdLeafHeader));
						if (elems[j] == selectedFace)
						{
							selectedNodes.Add(i);
							break;
						}
					}
	
				}
			}
		}
		else
		{
			if (GetAsyncKeyState(VK_CONTROL) != 0)
			{
				double model[16];
				double proj[16];
				int view[4];
				glGetDoublev(GL_MODELVIEW_MATRIX, model);
				glGetDoublev(GL_PROJECTION_MATRIX, proj);
				glGetIntegerv(GL_VIEWPORT, view);
				double dx, dy, dz;
				gluUnProject(args.X, GetClientHeight()-args.Y, 0.0f, model, proj, view, &dx, &dy, &dz);
				Matrix4 trans, trans2;
				Matrix4::RotationY(trans, viewTheta);
				Matrix4::RotationX(trans2, viewPhi);
				trans *= trans2;
				Vec3 dir(0.0f, 0.0f, -1.0f);
				Vec3 transDir;
				trans.Transform(transDir, dir);
				Vec3 pos = viewCenter;
				Vec3 dpos;
				Vec3::Scale(dpos, transDir, -1.0f * viewRadius);
				pos += dpos;
				Ray ray;
				ray.Origin = pos;
				ray.Direction = dir;
				ray.ReciprocalDirection.x = 1.0f/ray.Direction.x;
				ray.ReciprocalDirection.y = 1.0f/ray.Direction.y;
				ray.ReciprocalDirection.z = 1.0f/ray.Direction.z;
				ray.tMin = 0.001f;
				ray.tMax = FLT_MAX;
				SetTestRay(ray);
			}
			else
			{
				onAdjustView = true;
				lastX = args.X;
				lastY = args.Y;
				SetCapture(this->GetHandle());
			}
		}
	}

	void MainForm_MouseMove(Object * sender, MouseEventArgs args)
	{
		if (onAdjustView)
		{
			//if (args.Button == mbLeft)
			{
				viewPhi += (args.Y-lastY)/180.0f*PI;
				viewTheta += (lastX - args.X)/180.0f*PI;
				viewPhi = min(viewPhi, PI/2);
				viewPhi = max(viewPhi, -PI/2);
				lastX = args.X;
				lastY = args.Y;
			}
		}
	}

	void MainForm_MouseUp(Object * sender, MouseEventArgs args)
	{
		onAdjustView = false;
		ReleaseCapture();
	}
	
	void Mainform_Resized(Object * sender, EventArgs e)
	{
		SetProjection();
	}

	void MainForm_MouseWheel(Object * sender, MouseEventArgs e)
	{
		if (e.Delta > 0)
		{
			viewRadius -= 20.0f;
		}
		else
		{
			viewRadius += 20.0f;
		}
	}

	void MainForm_KeyPressed(Object * sender, KeyEventArgs e)
	{
		if (e.KeyCode == VK_RIGHT)
		{
			highlightStep ++;
		}
		else if (e.KeyCode == VK_LEFT)
		{
			highlightStep --;
		}
		if (tracedNodes.Count() > 0)
		{
			if (highlightStep < 0)
				highlightStep += tracedNodes.Count();
			highlightStep %= tracedNodes.Count();
			printf("highlighting node %d\n", tracedNodes[highlightStep]);
			PrintNodeInfo(mesh->GetKdTree(), tracedNodes[highlightStep]);
		}
		
	}

	// Trace Triangle
	inline bool TraceTriangle(DifferentialGeometry & inter, const MeshFace & face, int faceId, const Ray & ray, const RayDifferential & diffIn, float & t) const
	{
		//MeshTracingTimer(pred?2:1);
		const int mod3[] = {0,1,2,0,1};
		int u = mod3[face.ProjectionAxis + 1];
		int v = mod3[face.ProjectionAxis + 2];
		float invNdotD = 1.0f / (face.PlaneU * ray.Direction[u] + face.PlaneV * ray.Direction[v] + ray.Direction[face.ProjectionAxis]);
		float tplane = -(face.PlaneU * ray.Origin[u] + face.PlaneV * ray.Origin[v] + ray.Origin[face.ProjectionAxis] + face.PlaneD) * invNdotD;
		if (tplane < ray.tMin/*-Epsilon*/ || tplane > ray.tMax/*+Epsilon*/)
			return false;
		float hitU = ray.Origin[u] + ray.Direction[u] * tplane;
		float hitV = ray.Origin[v] + ray.Direction[v] * tplane;
			
		float beta = face.K_beta_u * hitU + face.K_beta_v * hitV + face.K_beta_d;
		if (beta < -Epsilon)
			return false;
		float gamma = face.K_gamma_u * hitU + face.K_gamma_v * hitV + face.K_gamma_d;
		if (gamma < -Epsilon)
			return false;
		if (beta + gamma > 1.0f + Epsilon)
			return false;
		inter.Position[u] = hitU;
		inter.Position[v] = hitV;
		inter.Position[face.ProjectionAxis] = ray.Origin[face.ProjectionAxis] + ray.Direction[face.ProjectionAxis] * tplane;
		//float alpha = face.K_alpha_u * hitU + face.K_alpha_v * hitV + face.K_alpha_d;
		float alpha = 1.0f - beta - gamma;
		Vec3 tmp;
		Vec3::Scale(inter.Normal, mesh->Normals[face.Normal1], alpha);
		Vec3::Scale(tmp, mesh->Normals[face.Normal2], beta);
		inter.Normal += tmp;
		Vec3::Scale(tmp, mesh->Normals[face.Normal3], gamma);
		inter.Normal += tmp;

		Vec2 tmpVec2;
		if (face.TexCoord1 != -1)
		{
			inter.UV.x = mesh->TexCoords[face.TexCoord1].x * alpha;
			inter.UV.y = mesh->TexCoords[face.TexCoord1].y * alpha;
			inter.UV.x += mesh->TexCoords[face.TexCoord2].x * beta;
			inter.UV.y += mesh->TexCoords[face.TexCoord2].y * beta;
			inter.UV.x += mesh->TexCoords[face.TexCoord3].x * gamma;
			inter.UV.y += mesh->TexCoords[face.TexCoord3].y * gamma;
		}
		inter.dNdxContinuation = faceId;
		t = tplane;
		return true;
	}

	void PrintNodeInfo(KdTree<MeshFace> * tree, int nodeId)
	{
		printf("  bbox: %f, %f, %f -- %f, %f, %f\n", nodeBounds[nodeId].xMin,nodeBounds[nodeId].yMin, nodeBounds[nodeId].zMin, nodeBounds[nodeId].xMax, nodeBounds[nodeId].yMax, nodeBounds[nodeId].zMax);
		if (tree->Nodes[nodeId].GetIsLeaf())
		{
			printf("  #triangles: %d\n", tree->Nodes[nodeId].GetElementCount());
			printf("  intersections:\n");

			highlightFaces.clear();
			int * elems = (int *)((char*)tree->Leaves.Buffer() + tree->Nodes[nodeId].LeafPtr + sizeof(KdLeafHeader));
			for (int i = 0; i<tree->Nodes[nodeId].GetElementCount(); i++)
			{
				DifferentialGeometry inter;
				RayDifferential diffIn;
				float t;
				highlightFaces.insert(elems[i]);
				bool succ = TraceTriangle(inter, tree->Elements[elems[i]], elems[i], testRay, diffIn, t); 
				if (succ)
				{
					printf("    hit pos(%f, %f, %f)\n", inter.Position.x, inter.Position.y, inter.Position.z);
					printf("        normal(%f, %f, %f)\n", inter.Normal.x, inter.Normal.y, inter.Normal.z);
					printf("        t = %f\n", t);
				}
			}
		}
	}
public:
	MainForm()
	{
		SetText(L"KdTree Visualizer");
		SetClientWidth(800);
		SetClientHeight(600);
		this->OnMouseDown.Bind(this, &MainForm::MainForm_MouseDown);
		this->OnMouseMove.Bind(this, &MainForm::MainForm_MouseMove);
		this->OnMouseUp.Bind(this, &MainForm::MainForm_MouseUp);
		this->OnResized.Bind(this, &MainForm::Mainform_Resized);
		this->OnMouseWheel.Bind(this, &MainForm::MainForm_MouseWheel);
		this->OnKeyDown.Bind(this, &MainForm::MainForm_KeyPressed);
		timer = new Timer();
		timer->Interval = 20;
		timer->OnTick.Bind(this, &MainForm::RefreshTimer_OnTick);
		timer->StartTimer();

		system = new RayTraceSystem();
		InitModel();
		InitRendering();

		onAdjustView = false;
		selectedFace = -1;
		highlightStep = 0;
	}

	~MainForm()
	{
		delete system;
	}

	void UpdateViewTransform()
	{
		Matrix4 trans, trans2;
	
		Matrix4::RotationY(trans, viewTheta);
		Matrix4::RotationX(trans2, viewPhi);
		trans *= trans2;
		Vec3 dir(0.0f, 0.0f, -1.0f);
		Vec3 transDir;
		trans.Transform(transDir, dir);
		Vec3 pos = viewCenter;
		Vec3 dpos;
		Vec3::Scale(dpos, transDir, -1.0f * viewRadius);
		pos += dpos;
		Matrix4::Translation(trans2, pos.x, pos.y, pos.z);
		trans.LeftMultiply(trans2);
		
		Matrix4 inv;
		trans.Inverse(inv);
		glLoadMatrixf(inv.values);
	}

	void DrawCube(BBox & bounds)
	{
		glBegin(GL_QUADS);
		glVertex3f(bounds.xMin, bounds.yMin, bounds.zMin);
		glVertex3f(bounds.xMax, bounds.yMin, bounds.zMin);
		glVertex3f(bounds.xMax, bounds.yMin, bounds.zMax);
		glVertex3f(bounds.xMin, bounds.yMin, bounds.zMax);

		glVertex3f(bounds.xMin, bounds.yMax, bounds.zMin);
		glVertex3f(bounds.xMax, bounds.yMax, bounds.zMin);
		glVertex3f(bounds.xMax, bounds.yMax, bounds.zMax);
		glVertex3f(bounds.xMin, bounds.yMax, bounds.zMax);

		glVertex3f(bounds.xMin, bounds.yMin, bounds.zMin);
		glVertex3f(bounds.xMax, bounds.yMin, bounds.zMin);
		glVertex3f(bounds.xMax, bounds.yMax, bounds.zMin);
		glVertex3f(bounds.xMin, bounds.yMax, bounds.zMin);

		glVertex3f(bounds.xMin, bounds.yMin, bounds.zMax);
		glVertex3f(bounds.xMax, bounds.yMin, bounds.zMax);
		glVertex3f(bounds.xMax, bounds.yMax, bounds.zMax);
		glVertex3f(bounds.xMin, bounds.yMax, bounds.zMax);

		glVertex3f(bounds.xMin, bounds.yMin, bounds.zMin);
		glVertex3f(bounds.xMin, bounds.yMax, bounds.zMin);
		glVertex3f(bounds.xMin, bounds.yMax, bounds.zMax);
		glVertex3f(bounds.xMin, bounds.yMin, bounds.zMax);

		glVertex3f(bounds.xMax, bounds.yMin, bounds.zMin);
		glVertex3f(bounds.xMax, bounds.yMax, bounds.zMin);
		glVertex3f(bounds.xMax, bounds.yMax, bounds.zMax);
		glVertex3f(bounds.xMax, bounds.yMin, bounds.zMax);
		glEnd();
	}

	void DrawBounds(BBox & bounds)
	{
		
		DrawCube(bounds);
	}

	void RefreshTimer_OnTick(Object * sender, EventArgs e)
	{
		BeginGLDrawing();
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
		UpdateViewTransform();
		glColor3ub(255,255,255);
		DrawModel();
		
		glBegin(GL_LINES);
		glColor4ub(255, 255, 0, 255);
		glVertex3f(testRay.Origin.x, testRay.Origin.y, testRay.Origin.z);
		glVertex3f(testRay.Origin.x + testRay.Direction.x*testRay.tMin, 
			testRay.Origin.y + testRay.Direction.y*testRay.tMin, 
			testRay.Origin.z + testRay.Direction.z*testRay.tMin);
		glColor4ub(0, 255, 0, 255);
		glVertex3f(testRay.Origin.x + testRay.Direction.x*testRay.tMin, 
			testRay.Origin.y + testRay.Direction.y*testRay.tMin, 
			testRay.Origin.z + testRay.Direction.z*testRay.tMin);
		glVertex3f(testRay.Origin.x + testRay.Direction.x*testRay.tMax, 
			testRay.Origin.y + testRay.Direction.y*testRay.tMax, 
			testRay.Origin.z + testRay.Direction.z*testRay.tMax);
		glEnd();
		
		glColor4ub(255, 127, 127, 127);

		for (int i = 0; i<selectedNodes.Count(); i++)
		{
			DrawBounds(nodeBounds[selectedNodes[i]]);
		}
		glEnable(GL_POLYGON_OFFSET_FILL);
		glPolygonOffset(-0.1f, 0.1f);
		for (int i = 0; i<tracedNodes.Count(); i++)
		{
			if (highlightStep == i)
			{
				glColor4ub(255, 127, 0, 127);
				glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);
			}
			else if (!tree->Nodes[tracedNodes[i]].GetIsLeaf())
			{
				glColor4ub(255, 0, 0, 127);
				glPolygonMode(GL_FRONT_AND_BACK, GL_LINE);
			}
			else
			{
				glColor4ub(0, 0, 255, 127);
				glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);
			}
			DrawBounds(nodeBounds[tracedNodes[i]]);
		}
		glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);
		glDisable(GL_POLYGON_OFFSET_FILL);
		glPolygonOffset(GL_POLYGON_OFFSET_FILL, 0.0f);
		EndGLDrawing();
	}
};

int __stdcall _tWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPTSTR lpCmdLine, int nCmdShow) 
{
	Application::Init(hInstance, hPrevInstance, lpCmdLine, nCmdShow);
	MainForm * form = new MainForm();
	Application::Run(form);
	Application::Dispose();
	return 0;
}

int _tmain(int argc, TCHAR* argv[])
{
	return _tWinMain((HINSTANCE)GetModuleHandle(NULL), NULL, GetCommandLine(), 0);
}

